package org.dromara.web.plugins.provider.imageImpl;

import com.alibaba.fastjson.JSONArray;
import com.alibaba.fastjson.JSONObject;
import okhttp3.*;
import org.dromara.common.core.exception.ServiceException;
import org.dromara.common.mongo.ChatMessageMongoDao;
import org.dromara.common.mongo.SessionMongoDao;
import org.dromara.dao.business.domain.Model;
import org.dromara.dao.business.domain.ModelParam;
import org.dromara.web.plugins.provider.TextToImageProvider;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;
import org.springframework.stereotype.Component;

import java.util.HashMap;
import java.util.Map;


@Component
public class DalleImageProvider implements TextToImageProvider {

    public static final String MODEL_TYPE = "openai";

    private static final String URL = "https://api.openai.com/v1/images/generations";
    private static final String API_KEY = "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";


    @Autowired
    OkHttpClient okHttpClient;
    @Autowired
    ThreadPoolTaskExecutor threadPoolTaskExecutor;
    @Autowired
    SessionMongoDao sessionMongoDao;
    @Autowired
    ChatMessageMongoDao chatMessageMongoDao;

    public static final Map<String, String> sizeMap = new HashMap<>();
    static {
        sizeMap.put("4:7", "1024x1792");
        sizeMap.put("1:1", "1024x1024");
        sizeMap.put("7:4", "1792x1024");
    }

    @Override
    public String generateImage(Model model, String prompt, ModelParam modelParam, String scale) {


        FormBody.Builder builder = new FormBody.Builder();
        builder.add("model", model.getValue());
        builder.add("prompt", prompt);
        builder.add("style", modelParam.getPromt());
        builder.add("size", sizeMap.get(scale));
        builder.add("response_format", "b64_json");



        RequestBody requestBody = builder.build();

        Request request = new Request.Builder()
            .url(URL)
            .header("Authorization", "Bearer " + API_KEY)
            .post(requestBody)
            .build();


        try {
            Response response = okHttpClient.newCall(request).execute();
            if (response.isSuccessful()) {
                JSONObject jsonObject = JSONObject.parseObject(response.body().string());
                JSONArray jsonArray = jsonObject.getJSONArray("data");
                if (jsonArray != null && jsonArray.size() > 0) {
                    JSONObject object = jsonArray.getJSONObject(0);
                    return "data:image/jpeg;base64," + object.getString("url");
                }
            }

        } catch (Exception e) {
        }
        throw new ServiceException("服务异常");
    }
}
